# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: wooyaggo

from recordclass import recordclass

from coin.base.book import BookBuilderBase
from coin.base.book.book_builder_base import MutablePriceQty

OrderEntry = recordclass('OrderEntry', ['orderID', 'price', 'side', 'leavesQty', 'symbol'])


def new_order_entry(arg):
  if isinstance(arg, OrderEntry):
    dic = arg._asdict()
  elif isinstance(arg, dict):
    dic = arg
  else:
    raise NotImplementedError('Not supported argument: ' + repr(arg))

  for field in OrderEntry._fields:
    if field not in dic:
      dic[field] = None
  return OrderEntry(**dic)


class LiquidationBookBuilder(BookBuilderBase):
  def __init__(self, symbol):
    super().__init__("Bitmex", symbol)
    self._id_map = {}

  def delete(self, id: str):
    if id not in self._id_map:
      return

    order = self._id_map[id]
    price = order.price
    side = order.side
    qty = order.leavesQty

    del self._id_map[id]
    if side == 'Sell':
      self._ask_price_map.clear_cache()
      self._ask_price_map[price].qty -= qty
      assert self._ask_price_map[price].qty >= 0, self._ask_price_map[price]

      if self._ask_price_map[price].qty == 0:
        del self._ask_price_map[price]
    elif side == 'Buy':
      self._bid_price_map.clear_cache()
      self._bid_price_map[price].qty -= qty
      assert self._bid_price_map[price].qty >= 0, self._bid_price_map[price]

      if self._bid_price_map[price].qty == 0:
        del self._bid_price_map[price]
    else:
      raise ValueError('Invalid side: %s' % side)

  def update(self, id: str, amend_order: OrderEntry):
    if id not in self._id_map:
      return
    old_order = self._id_map[id]

    self.delete(id)

    new_order = new_order_entry(old_order)
    if amend_order.price is not None:
      new_order.price = amend_order.price
    elif amend_order.leavesQty is not None:
      new_order.leavesQty = amend_order.leavesQty
    elif amend_order.side is not None:
      new_order.side = amend_order.side

    self.insert(new_order)

  def insert(self, order: OrderEntry):
    id = order.orderID
    price = order.price
    qty = order.leavesQty
    side = order.side

    assert id not in self._id_map, 'ID not exist in map: ' + id
    assert price > 0
    assert qty > 0, 'Size must be greater than 0.'

    if side == 'Sell':
      if price in self._ask_price_map:
        self._ask_price_map.clear_cache()
        self._ask_price_map[price].qty += qty
      else:
        self._ask_price_map[price] = MutablePriceQty(price, qty)

    elif side == 'Buy':
      if price in self._bid_price_map:
        self._bid_price_map.clear_cache()
        self._bid_price_map[price].qty += qty
      else:
        self._bid_price_map[price] = MutablePriceQty(price, qty)

    else:
      raise ValueError('Invalid side: %s' % side)

    self._id_map[id] = new_order_entry(order)
